import os, cv2
from glob import glob
import SimpleITK as sitk
import numpy as np

def turn2DTo3DV1(path_folder='', path_save='demo.nii'):
    imgs = glob(path_folder+'/*.jpg')
    assert(len(imgs)>0, "There are no images in folder:"+path_folder)

    imgs = [cv2.imread(img, cv2.IMREAD_GRAYSCALE) for img in imgs]
    # imgs = [sitk.ReadImage(img) for img in imgs]
    # imgs = [sitk.GetArrayFromImage(img) for img in imgs]
    print(imgs[0].shape)

    imgs = np.stack(imgs, axis=0)
    imgs = sitk.GetImageFromArray(imgs)
    if not path_save.endswith('.nii'):
        path_save += '.nii'
    sitk.WriteImage(imgs, path_save)


def turn2DTo3DV2(path_folder='', path_save='demo.nii'):
    imgs = []
    for root, dirs, files in os.walk(path_folder):
        # 遍历文件
        for f in files:
            print('file:', os.path.join(root, f))
            if f.split('.')[-1].lower() in ['jpg', 'jpeg', 'bmp', 'png', 'gif', 'tif']:
            # if f.endswith('.jpg') or f.endswith('.png') or f.endswith('.bmp') or f.endswith('.jpeg'):
                imgs.append(os.path.join(root, f))

    # imgs = glob(path_folder+'/*.jpg')
    assert(len(imgs)>0, "There are no images in folder:"+path_folder)
    # print(imgs[0])
    imgs = [sitk.ReadImage(img) for img in imgs]
    imgs = [sitk.GetArrayFromImage(img) for img in imgs]
    imgs = [np.mean(img, axis=-1) for img in imgs]
    # print(imgs[0].shape)

    imgs = np.stack(imgs, axis=0)
    imgs = sitk.GetImageFromArray(imgs)
    if not path_save.endswith('.nii'):
        path_save += '.nii'
    sitk.WriteImage(imgs, path_save)

'''
环境：Python3(opencv, SimpleITK, numpy, glob)
把文件夹path_folder下的所有2维图片打包成单个三维格式
输入参数：
    path_folder:要打包的二维图片文件夹路径
    path_save:要把打包好的nii文件保存到的路径
'''
if __name__ == '__main__':
    folder = r'G:\Desktop\USMITK\data_image'
    # turn2DTo3DV1(folder, 'demo.nii')
    turn2DTo3DV2(folder, 'demo2.nii')